import numpy as np
import torch
import time
import importlib
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  #CPU MODEL ONLY
import importlib
import argparse

def main(net_name, model_path, input_shape, num_classes):

    module = importlib.import_module("nets.%s.model" % net_name)
    model = module.get_seg_model(num_classes, aux=True)
    model = model.to("cpu")
    model.eval()

    # qmodel = torch.quantization.convert(model)

    dummy_input = torch.rand(1, 3, input_shape[1], input_shape[0])
    torch.save(model, "preview/%s.pth" % net_name)


    try:
        print("start to convert jit...")
        trace_model = torch.jit.trace(model, dummy_input)
        trace_model.save("preview/%s.pt" % net_name)
        print("convert jit sucessful")
    except Exception as e:
        print("!!!!!!!!!!!convert jit failed\n")


    print('-------------------------------------------------')
    try:
        print("start to convert onnx...")
        torch.onnx.export(model, dummy_input, "preview/%s.onnx"%net_name)#, opset_version=11)
        print("convert onnx sucessful")
    except Exception as e:
        print("!!!!!!!!!!!!convert onnx failed")



if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--net_name", required=True)
    parser.add_argument("--model_path", type=str, default="")
    parser.add_argument('--input_shape', type=int, nargs="+", default=(224, 224, 3))
    parser.add_argument('--num_classes', type=int, default=2)
    args = parser.parse_args()
    main(args.net_name, args.model_path, args.input_shape, args.num_classes)

